package com.hc.java.pkcs.aes;

import com.hc.java.hash.Base64Utils;
import com.hc.java.hash.Hash;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Base64;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 工作模式:
 * 1.ECB(Electronic Codebook) 电子密码本模式, 加密和解密可并行, 明文长度必须是8字节的倍数(所以必须有填充方式), 不能设置iv
 * 2.CBC(Cipher Block Chaining) 密码分组链接模式, 加密不可并行, 解密可并行, 明文长度必须填充至8字节的倍数(所以必须有填充方式), 必须设置iv
 * 3.CFB(Cipher Feedback) 密文反馈模式, 加密不可并行, 解密可并行, 明文长度无要求, 必须设置iv
 * 4.OFB(Output Feedback) 输出反馈模式, 加密部分可并行, 解密可并行, 明文长度无要求, 必须设置iv
 * 5.CTR(Counter) 计数器模式, 加密和解密可并行, 明文长度无要求, 必须设置iv
 * <p>
 * 填充方式:
 * 1.NoPadding 无填充
 * 2.PKCS5Padding 它的块大小固定为8字节
 * 3.PKCS7Padding 它的块大小不固定, 块大小的范围是1-255字节
 */
public class Aes {

    private final static byte[] EMPTY_BYTES = new byte[0];

    //不推荐
    public static String AES_EBC_PKCS5PADDING = "AES/ECB/PKCS5Padding";

    //推荐
    public static String AES_CBC_PKCS5PADDING = "AES/CBC/PKCS5Padding";

    //看情况使用
    public static String AES_CFB_NOPADDING = "AES/CFB/NoPadding";
    public static String AES_CFB_PKCS5PADDING = "AES/CFB/PKCS5Padding";

    //看情况使用
    public static String AES_OFB_NOPADDING = "AES/OFB/NoPadding";
    public static String AES_OFB_PKCS5PADDING = "AES/OFB/PKCS5Padding";

    //推荐
    public static String AES_CTR_NOPADDING = "AES/CTR/NoPadding";

    //允许的算法
    private static final Set<String> allowedAlgorithms = new HashSet<>();

    static {
        allowedAlgorithms.add(AES_EBC_PKCS5PADDING);
        allowedAlgorithms.add(AES_CBC_PKCS5PADDING);
        allowedAlgorithms.add(AES_CFB_NOPADDING);
        allowedAlgorithms.add(AES_CFB_PKCS5PADDING);
        allowedAlgorithms.add(AES_OFB_NOPADDING);
        allowedAlgorithms.add(AES_OFB_PKCS5PADDING);
        allowedAlgorithms.add(AES_CTR_NOPADDING);
    }


    //算法
    private final String algorithm;

    //密钥
    private final SecretKey secretKey;

    //初始化向量
    private final IvParameterSpec iv;


    //加密器
    private final Cipher encryptCipher;
    //解密器
    private final Cipher decryptCipher;

    /**
     * 构造函数
     *
     * @param algorithm 算法
     * @param secretKey 密钥
     * @param iv        初始化向量
     */
    private Aes(String algorithm, SecretKey secretKey, IvParameterSpec iv) {
        this.algorithm = algorithm;
        this.secretKey = secretKey;
        this.iv = iv;

        this.encryptCipher = createEncryptCipher(algorithm, secretKey, iv);
        this.decryptCipher = createDecryptCipher(algorithm, secretKey, iv);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后的数据
     */
    public byte[] encryptToBytes(byte[] data) {
        if (null == data) {
            return EMPTY_BYTES;
        }

        try {
            return encryptCipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException("unable to encrypt data", e);
        }
    }

    /**
     * 加密
     *
     * @param data    待加密数据
     * @param charset data的字符集
     * @return 加密后的数据
     */
    public byte[] encryptToBytes(String data, Charset charset) {
        if (null == data) {
            return EMPTY_BYTES;
        }

        if (charset == null) {
            charset = StandardCharsets.UTF_8;
        }

        byte[] bytes = data.getBytes(charset);
        return encryptToBytes(bytes);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后的数据
     */
    public byte[] encryptToBytes(String data) {
        return encryptToBytes(data, StandardCharsets.UTF_8);
    }

    /**
     * 从base64字符串加密
     *
     * @param base64Data 使用base64编码的待加密数据
     * @return 加密后的数据
     */
    public byte[] encryptToBytesFromBase64(String base64Data) {
        if (null == base64Data || base64Data.isEmpty()) {
            return EMPTY_BYTES;
        }

        byte[] utf8Bytes = base64Data.getBytes(StandardCharsets.UTF_8);
        byte[] decoded = Base64.getDecoder().decode(utf8Bytes);
        return encryptToBytes(decoded);
    }

    /**
     * 从hex字符串加密
     *
     * @param hexData 用hex表示的待加密数据
     * @return 加密后的数据
     */
    public byte[] encryptToBytesFromHex(String hexData) {
        byte[] decoded = Hash.toHexBytes(hexData);
        return encryptToBytes(decoded);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后并用base64编码后的数据
     */
    public String encryptToBase64(byte[] data) {
        byte[] encrypted = encryptToBytes(data);
        return Base64Utils.encodeToString(encrypted);
    }

    /**
     * 加密
     *
     * @param data    待加密数据
     * @param charset data的字符集
     * @return 加密后并用base64编码后的数据
     */
    public String encryptToBase64(String data, Charset charset) {
        byte[] encrypted = encryptToBytes(data, charset);
        return Base64Utils.encodeToString(encrypted);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后并用base64编码后的数据
     */
    public String encryptToBase64(String data) {
        return encryptToBase64(data, StandardCharsets.UTF_8);
    }

    /**
     * 加密
     *
     * @param base64Data 用base64编码的待加密数据
     * @return 加密后并用base64编码后的数据
     */
    public String encryptToBase64FromBase64(String base64Data) {
        if (null == base64Data || base64Data.isEmpty()) {
            return "";
        }

        byte[] decoded = Base64Utils.decodeToBytes(base64Data);
        byte[] encrypted = encryptToBytes(decoded);
        return Base64Utils.encodeToString(encrypted);
    }

    /**
     * 加密
     *
     * @param hexData 用hex表示的待加密数据
     * @return 加密后并用base64编码后的数据
     */
    public String encryptToBase64FromHex(String hexData) {
        byte[] decoded = Hash.toHexBytes(hexData);
        byte[] encrypted = encryptToBytes(decoded);
        return Base64Utils.encodeToString(encrypted);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后并用hex表示的数据
     */
    public String encryptToHex(byte[] data) {
        byte[] encrypted = encryptToBytes(data);
        return Hash.toHexLowerCase(encrypted);
    }

    /**
     * 加密
     *
     * @param data    待加密数据
     * @param charset data的字符集
     * @return 加密后并用hex表示的数据
     */
    public String encryptToHex(String data, Charset charset) {
        byte[] encrypted = encryptToBytes(data, charset);
        return Hash.toHexLowerCase(encrypted);
    }

    /**
     * 加密
     *
     * @param data 待加密数据
     * @return 加密后并用hex表示的数据
     */
    public String encryptToHex(String data) {
        return encryptToHex(data, StandardCharsets.UTF_8);
    }

    /**
     * 加密
     *
     * @param base64Data 用base64编码的待加密数据
     * @return 加密后并用hex表示的数据
     */
    public String encryptToHexFromBase64(String base64Data) {
        byte[] decoded = Base64Utils.decodeToBytes(base64Data);
        byte[] encrypted = encryptToBytes(decoded);
        return Hash.toHexLowerCase(encrypted);
    }

    /**
     * 加密
     *
     * @param hexData 用hex表示的待加密数据
     * @return 加密后并用hex表示的数据
     */
    public String encryptToHexFromHex(String hexData) {
        byte[] decoded = Hash.toHexBytes(hexData);
        byte[] encrypted = encryptToBytes(decoded);
        return Hash.toHexLowerCase(encrypted);
    }

    /**
     * 解密
     *
     * @param data 待解密数据
     * @return 解密后的数据
     */
    public byte[] decryptToBytes(byte[] data) {
        if (null == data) {
            return EMPTY_BYTES;
        }

        try {
            return decryptCipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException("unable to decrypt data", e);
        }
    }

    /**
     * 解密
     *
     * @param data    待解密数据
     * @param charset data的字符集
     * @return 解密后的数据
     */
    public byte[] decryptToBytes(String data, Charset charset) {
        if (null == data) {
            return EMPTY_BYTES;
        }

        if (charset == null) {
            charset = StandardCharsets.UTF_8;
        }

        byte[] bytes = data.getBytes(charset);
        return decryptToBytes(bytes);
    }

    /**
     * 解密
     *
     * @param data 待解密数据
     * @return 解密后的数据
     */
    public byte[] decryptToBytes(String data) {
        return decryptToBytes(data, StandardCharsets.UTF_8);
    }

    /**
     * 从base64字符串编码的数据解密
     *
     * @param base64Data 用base64编码的待解密数据
     * @return 解密后的数据
     */
    public byte[] decryptToBytesFromBase64(String base64Data) {
        if (null == base64Data) {
            return EMPTY_BYTES;
        }

        byte[] utf8Bytes = base64Data.getBytes(StandardCharsets.UTF_8);
        byte[] decoded = Base64.getDecoder().decode(utf8Bytes);
        return decryptToBytes(decoded);
    }

    /**
     * 从hex字符串表示的数据解密
     *
     * @param hexData 用hex表示的待解密数据
     * @return 解密后的数据
     */
    public byte[] decryptToBytesFromHex(String hexData) {
        byte[] decoded = Hash.toHexBytes(hexData);
        return decryptToBytes(decoded);
    }

    /**
     * 解密
     *
     * @param data 待解密数据
     * @return 解密后的数据
     */
    public String decryptToString(byte[] data) {
        byte[] decrypted = decryptToBytes(data);
        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /**
     * 解密
     *
     * @param data    待解密数据
     * @param charset data的字符集
     * @return 解密后的数据
     */
    public String decryptToString(String data, Charset charset) {
        byte[] decrypted = decryptToBytes(data, charset);
        return new String(decrypted, charset);
    }

    /**
     * 解密
     *
     * @param data 待解密数据
     * @return 解密后的数据
     */
    public String decryptToString(String data) {
        return decryptToString(data, StandardCharsets.UTF_8);
    }

    /**
     * 从base64字符串编码的数据解密
     *
     * @param base64Data 用base64编码的待解密数据
     * @return 解密后的数据
     */
    public String decryptToStringFromBase64(String base64Data) {
        byte[] decrypted = decryptToBytesFromBase64(base64Data);
        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /**
     * 从hex字符串表示的数据解密
     *
     * @param hexData 用hex表示的待解密数据
     * @return 解密后的数据
     */
    public String decryptToStringFromHex(String hexData) {
        byte[] decrypted = decryptToBytesFromHex(hexData);
        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /**
     * 生成新的aes密钥
     *
     * @param size 密钥长度, 可选值为: 128, 192, 256
     * @return 新的AES密钥
     */
    public static SecretKey newKey(int size) {
        if (size != 128 && size != 192 && size != 256) {
            throw new IllegalArgumentException("invalid key size, only 128, 192, and 256 bits are supported");
        }

        try {
            KeyGenerator generator = KeyGenerator.getInstance("AES");
            generator.init(size);
            return generator.generateKey();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("unable to generate secret key", e);
        }
    }

    /**
     * 把byte数组转换为aes的密钥
     *
     * @param key 密钥字节数组
     * @return 密钥
     */
    public static SecretKey toSecretKey(byte[] key) {
        if (null == key) {
            throw new IllegalArgumentException("key cannot be null");
        }

        if (key.length != 16 && key.length != 24 && key.length != 32) {
            throw new IllegalArgumentException("invalid key size, only 128, 192, and 256 bits are supported");
        }

        return new SecretKeySpec(key, "AES");
    }

    /**
     * 把字符串转换为aes的密钥
     *
     * @param key 密钥字符串
     * @return 密钥
     */
    public static SecretKey toSecretKey(String key) {
        if (null == key) {
            throw new IllegalArgumentException("key cannot be null");
        }

        byte[] decodedKey = key.getBytes(StandardCharsets.UTF_8);
        return toSecretKey(decodedKey);
    }

    /**
     * 把Base64字符串转换为aes的密钥
     *
     * @param base64Key 密钥Base64字符串
     * @return 密钥
     */
    public static SecretKey toSecretKeyFromBase64(String base64Key) {
        if (null == base64Key) {
            throw new IllegalArgumentException("key cannot be null");
        }

        byte[] utf8Bytes = base64Key.getBytes(StandardCharsets.UTF_8);
        byte[] decodedKey = Base64.getDecoder().decode(utf8Bytes);
        return toSecretKey(decodedKey);
    }

    /**
     * 把hex字符串转换为aes的密钥
     *
     * @param hexKey 密钥hex字符串
     * @return 密钥
     */
    public static SecretKey toSecretKeyFromHex(String hexKey) {
        if (null == hexKey) {
            throw new IllegalArgumentException("key cannot be null");
        }

        byte[] decodedKey = Hash.toHexBytes(hexKey);
        return toSecretKey(decodedKey);
    }

    /**
     * 生成新的aes初始化向量
     *
     * @param size 向量长度, 可选值为: 128, 192, 256
     * @return 新的aes初始化向量
     */
    public static IvParameterSpec newIv(int size) {
        ThreadLocalRandom random = ThreadLocalRandom.current();

        if (size == 128) {
            byte[] iv = new byte[16];
            random.nextBytes(iv);
            return new IvParameterSpec(iv);
        } else if (size == 192) {
            byte[] iv = new byte[24];
            random.nextBytes(iv);
            return new IvParameterSpec(iv);
        } else if (size == 256) {
            byte[] iv = new byte[32];
            random.nextBytes(iv);
            return new IvParameterSpec(iv);
        } else {
            throw new IllegalArgumentException("invalid key size, only 128, 192, and 256 bits are supported");
        }
    }

    /**
     * 把byte数组转换为aes的初始化向量
     *
     * @param iv 向量字节数组
     * @return 向量
     */
    public static IvParameterSpec toIv(byte[] iv) {
        if (null == iv) {
            throw new IllegalArgumentException("iv cannot be null");
        }

        if (iv.length != 16 && iv.length != 24 && iv.length != 32) {
            throw new IllegalArgumentException("invalid iv size, only 128, 192, and 256 bits are supported");
        }

        return new IvParameterSpec(iv);
    }

    /**
     * 把字符串转换为aes的初始化向量
     *
     * @param iv 向量字符串
     * @return 向量
     */
    public static IvParameterSpec toIv(String iv) {
        if (null == iv) {
            throw new IllegalArgumentException("iv cannot be null");
        }

        byte[] decodedIv = iv.getBytes(StandardCharsets.UTF_8);
        return toIv(decodedIv);
    }

    /**
     * 把Base64字符串转换为aes的初始化向量
     *
     * @param base64Iv 向量Base64字符串
     * @return 向量
     */
    public static IvParameterSpec toIvFromBase64(String base64Iv) {
        if (null == base64Iv) {
            throw new IllegalArgumentException("iv cannot be null");
        }

        byte[] utf8Bytes = base64Iv.getBytes(StandardCharsets.UTF_8);
        byte[] decodedIv = Base64.getDecoder().decode(utf8Bytes);
        return toIv(decodedIv);
    }

    /**
     * 把hex字符串转换为aes的初始化向量
     *
     * @param hexIv 向量hex字符串
     * @return 向量
     */
    public static IvParameterSpec toIvFromHex(String hexIv) {
        if (null == hexIv) {
            throw new IllegalArgumentException("iv cannot be null");
        }

        byte[] decodedIv = Hash.toHexBytes(hexIv);
        return toIv(decodedIv);
    }

    /**
     * 创建aes实例
     *
     * @param algorithm 算法
     * @param key       密钥
     * @param iv        初始化向量
     * @return aes实例
     */
    public static Aes of(String algorithm, SecretKey key, IvParameterSpec iv) {
        return new Aes(algorithm, key, iv);
    }

    /**
     * 创建aes实例
     *
     * @param algorithm 算法
     * @param key       密钥
     * @param iv        初始化向量
     * @return aes实例
     */
    public static Aes of(String algorithm, byte[] key, byte[] iv) {
        return of(algorithm, toSecretKey(key), toIv(iv));
    }

    /**
     * 创建aes实例
     *
     * @param algorithm 算法
     * @param key       密钥
     * @param iv        初始化向量
     * @return aes实例
     */
    public static Aes of(String algorithm, String key, String iv) {
        return of(algorithm, toSecretKey(key), toIv(iv));
    }

    /**
     * 创建aes实例
     *
     * @param algorithm 算法
     * @param base64Key base64字符串表示的密钥
     * @param base64Iv  base64字符串表示的初始化向量
     * @return aes实例
     */
    public static Aes ofFromBase64(String algorithm, String base64Key, String base64Iv) {
        return of(algorithm, toSecretKeyFromBase64(base64Key), toIvFromBase64(base64Iv));
    }

    /**
     * 创建aes实例
     *
     * @param algorithm 算法
     * @param hexKey    hex字符串表示的密钥
     * @param hexIv     hex字符串表示的初始化向量
     * @return aes实例
     */
    public static Aes ofFromHex(String algorithm, String hexKey, String hexIv) {
        return of(algorithm, toSecretKeyFromHex(hexKey), toIvFromHex(hexIv));
    }

    /**
     * 创建并初始化cipher实例
     *
     * @param algorithm 算法
     * @param secretKey 密钥
     * @param iv        初始化向量
     * @return aes实例
     */
    private static Cipher createEncryptCipher(String algorithm, SecretKey secretKey, IvParameterSpec iv) {
        if (!allowedAlgorithms.contains(algorithm)) {
            throw new IllegalArgumentException("invalid algorithm: " + algorithm);
        }

        Cipher cipher;
        try {
            cipher = Cipher.getInstance(algorithm);
        } catch (Exception e) {
            throw new IllegalArgumentException("invalid algorithm: " + algorithm, e);
        }

        init(cipher, Cipher.ENCRYPT_MODE, secretKey, iv);
        return cipher;
    }

    /**
     * 创建并初始化cipher实例
     *
     * @param algorithm 算法
     * @param secretKey 密钥
     * @param iv        初始化向量
     * @return aes实例
     */
    private static Cipher createDecryptCipher(String algorithm, SecretKey secretKey, IvParameterSpec iv) {
        if (!allowedAlgorithms.contains(algorithm)) {
            throw new IllegalArgumentException("invalid algorithm: " + algorithm);
        }

        Cipher cipher;
        try {
            cipher = Cipher.getInstance(algorithm);
        } catch (Exception e) {
            throw new IllegalArgumentException("invalid algorithm: " + algorithm, e);
        }

        init(cipher, Cipher.DECRYPT_MODE, secretKey, iv);
        return cipher;
    }

    /**
     * 初始化cipher实例
     *
     * @param cipher        cipher实例
     * @param mode          模式
     * @param secretKey     密钥
     * @param parameterSpec 算法参数
     */
    private static void init(Cipher cipher, int mode, SecretKey secretKey, AlgorithmParameterSpec parameterSpec) {
        try {
            if (null != parameterSpec) {
                cipher.init(mode, secretKey, parameterSpec);
            } else {
                cipher.init(mode, secretKey);
            }
        } catch (InvalidKeyException e) {
            throw new IllegalArgumentException("unable to initialize due to invalid secret key", e);
        } catch (InvalidAlgorithmParameterException e) {
            throw new IllegalArgumentException("unable to initialize due to invalid decryption parameter spec", e);
        }
    }

    public String getAlgorithm() {
        return algorithm;
    }

    public SecretKey getSecretKey() {
        return secretKey;
    }

    public IvParameterSpec getIv() {
        return iv;
    }
}
